/*
 * Copyright (C) 2011 ViewlinkDC++ Team
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 */

// viewUpdate.cpp : Defines the entry point for the application.
//

#include "stdafx.h"

#include <Psapi.h>
#include <commctrl.h>
#include <shellapi.h>
#include <string>
#include "viewUpdate.h"
#include "SingleInstance.h"
#include "UpdateDlg.h"
#include "viewUpdateLang.h"

// #define ONE_VIEW_UPDATE

#define MAX_ARGS_COMMAND_LINE 5
#define NTDDI_VERSION NTDDI_LONGHORN
// #define _SSA_DEBUG_TESTING_
#define _MINIMAL_PID_DETECTION_

#define WIN32_LEAN_AND_MEAN		// Exclude rarely-used stuff from Windows headers

#define IMPLEMENT_VISTA_TOOLS
#include "VistaTools.cxx"

using namespace std;

#define MAX_LOADSTRING 100

// Global Variables:
HINSTANCE hInst;								// current instance
TCHAR szTitle[MAX_LOADSTRING];					// The title bar text
TCHAR szWindowClass[MAX_LOADSTRING];			// the main window class name


static bool is_wine()
{
	HMODULE module = GetModuleHandleA("ntdll.dll");
	if (!module)
		return false;
	return (GetProcAddress(module, "wine_server_call") != NULL);
}

bool CheckProcessID(wstring wProcessID, wstring fullName)
{
	bool retValue = false;
	DWORD procID = _wtoi(wProcessID.c_str());

	DWORD cbNeeded, cProcesses;

	DWORD countOfProc = 0;
	DWORD* aProcesses = NULL;


	do {

		countOfProc += 1024;

		if (aProcesses)
			delete[] aProcesses;
		aProcesses = NULL;

		aProcesses = new DWORD[countOfProc];

		cbNeeded = 0;
		
		if ( !EnumProcesses( aProcesses, countOfProc*sizeof(DWORD), &cbNeeded ) )
		{
			if (aProcesses)
				delete[] aProcesses;
			return true; // Problem, don't update!!!
		}

		cProcesses = cbNeeded / sizeof(DWORD);

	} while ( cProcesses == countOfProc);

#if defined(_SSA_DEBUG_TESTING_)
	wchar_t buff[256];
	wsprintf(buff, L"Searching PID = %d, Detected %d processes count", procID, cProcesses);
	MessageBox(0,buff, LAPP_CAPTION, MB_OK);
	wstring processes;
#endif

	for (DWORD i = 0; i< cProcesses; i++)
	{
#if defined(_SSA_DEBUG_TESTING_)
		{
		wchar_t buff[256];
		wsprintf(buff, L" Detected %d \t", aProcesses[i]);
		processes+=buff;
		}
#endif
		if (aProcesses[i] == procID)
		{

#if defined (_MINIMAL_PID_DETECTION_)
			retValue = true;

#if defined(_SSA_DEBUG_TESTING_)
		MessageBox(0,L"Process founded.", LAPP_CAPTION, MB_OK);
#endif

#else
#if defined(_SSA_DEBUG_TESTING_)
		MessageBox(0,L"Process founded. Try to open", LAPP_CAPTION, MB_OK);
#endif
			// Get Process Name
			HANDLE hProcess = OpenProcess( PROCESS_QUERY_INFORMATION |
				PROCESS_VM_READ,
				FALSE, procID );

			if (NULL != hProcess )
			{
#if defined(_SSA_DEBUG_TESTING_)
		MessageBox(0,L"Process Opened. Try to detect Imagefilename.", LAPP_CAPTION, MB_OK);
#endif

				TCHAR szProcessName[MAX_PATH] = TEXT("<unknown>");

				if (GetProcessImageFileName(hProcess, szProcessName, MAX_PATH))
				{
#if defined(_SSA_DEBUG_TESTING_)
				wstring processFound = L"Filename =";
				processFound+= szProcessName;
				MessageBox(0,processFound.c_str(), LAPP_CAPTION, MB_OK);
#endif
					// Remove driver name (C: for ex.)
					wstring localName = fullName;
					size_t lastValue = fullName.find_first_of(L":");
					if (lastValue && lastValue != wstring::npos && lastValue+1 < fullName.length())
						localName = fullName.substr( lastValue + 1);

#if defined(_SSA_DEBUG_TESTING_)
				processFound = L"LocalName =";
				processFound+= localName;
				processFound+= L" ProcName = ";
				processFound+= szProcessName;
				MessageBox(0,processFound.c_str(), LAPP_CAPTION, MB_OK);
#endif

					wstring procName = szProcessName;
					if (procName.find(localName) != wstring::npos)
						retValue = true;
				}else
				{
#if defined(_SSA_DEBUG_TESTING_)
				TCHAR szProcessError [256];
				wsprintf(szProcessError, L"GetProcessImageFileName error = %d", GetLastError());
				MessageBox(0,szProcessError, LAPP_CAPTION, MB_OK);
#endif
				}
				// [!] SSA - not work in x64 for x86 proc
				// DWORD err = 0;
				//HMODULE hMod = NULL;
				//DWORD cbNeeded;
				//if ( EnumProcessModules( hProcess, &hMod, sizeof(hMod), 
				//	&cbNeeded) )
				//{

				//	GetModuleFileNameEx( hProcess, hMod, szProcessName, 
				//		sizeof(szProcessName)/sizeof(TCHAR) );
				//	if ( fullName == szProcessName)
				//	{
				//		retValue = true;
				//	}
				//}else{
				//	err = GetLastError();
				//}
				
				CloseHandle( hProcess );
			}else
			{
#if defined(_SSA_DEBUG_TESTING_)
				TCHAR szProcessError [256];
				wsprintf(szProcessError, L"Process opening error = %d", GetLastError());
				MessageBox(0,szProcessError, LAPP_CAPTION, MB_OK);
#endif
			}

			if ( retValue )
				break;
#endif
		}
	}
	
	if (aProcesses)
		delete[] aProcesses;
	aProcesses = NULL;

#if defined(_SSA_DEBUG_TESTING_)
	MessageBox(0, processes.c_str(), LAPP_CAPTION, MB_OK);
#endif

	return retValue;
}

int APIENTRY _tWinMain(HINSTANCE hInstance,
	HINSTANCE hPrevInstance,
	LPTSTR    lpCmdLine,
	int       nCmdShow)
{
	UNREFERENCED_PARAMETER(hPrevInstance);

	
	HRESULT hr = CoInitializeEx(NULL, COINIT_MULTITHREADED);



	/// Check Command line 
	/// 1st parameter - Folder for ViewlinkDC++
	/// 2nd parameter - Folder with Update
	/// 3rd parameter - Need start After success = 0/1 (st)
	/// 4th parameter - ViewlinkDC++ ProcessID
	/// 5th parameter - currentState of Operation 
	///
	// 1 - wait for ViewlinkDC over
	// 2 - switch to AdminMode if need and copy from Update folder
	// 3 - switch to Usermode if need and load ViewlinkDC++
	// Command line for Start from View:
	// "C:\Program Files\2ViewlinkDC++\viewlinkdc.exe" "H:\Projects\viewlinkdc500\compiled\Update" 1 375 1
	//  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ^  ^  ^ 
	//           Path to viewlinkdc                          path to update folder               st id 1  

	if (lpCmdLine == NULL || wcslen(lpCmdLine) == 0)
		return 0;


#ifdef ONE_VIEW_UPDATE
	bool bIsUpdateRunning = false;
	int waitCountUpdate = 0;
	int waitCountMax = 10;
	do {
#ifndef _DEBUG
		SingleInstance dcUpapp(_T("{VIEWUPDDC-AEE8350A-B49A-4753-AB4B-E55479A48351}"));
#else
		SingleInstance dcUpapp(_T("{VIEWUPDDC-AEE8350A-B49A-4753-AB4B-E55479A48350}"));
#endif
		if ( bIsUpdateRunning  = ( dcUpapp.IsAnotherInstanceRunning() == TRUE) ){
			waitCountUpdate++;
			::Sleep(1000);
		}
		if (waitCountUpdate == waitCountMax)
			break;

	}while ( bIsUpdateRunning );


	if (bIsUpdateRunning)
		// Wait 10 sec to continue - else exit
		return 1;
#endif


	BOOL isElevated = FALSE;
	BOOL isVista =  is_wine() ? FALSE : IsVista();
	if (isVista == TRUE)
		IsElevated( &isElevated );

#if defined(_SSA_DEBUG_TESTING_)
	MessageBox(0, L"Checking argument lines", LAPP_CAPTION, MB_OK);
#endif
	// Parse line
	wstring arglines[MAX_ARGS_COMMAND_LINE];
	int iCount = 0;
	bool isCommas = false;
	for (UINT i = 0; i<wcslen(lpCmdLine);i++)
	{
		if (lpCmdLine[i] == L' ' && !isCommas)
		{
			iCount++;
		}else
		{
			if (lpCmdLine[i] == L'"')
				isCommas = !isCommas;
			arglines[iCount]+= lpCmdLine[i];
		}
	}

	wstring arglinesProcessed[MAX_ARGS_COMMAND_LINE];
	for (UINT i = 0; i<MAX_ARGS_COMMAND_LINE;i++)
	{
		// remove ""
		if (arglines[i].c_str()[0] == L'"' && arglines[i].c_str()[arglines[i].length()-1] == L'"' )
			arglinesProcessed[i] = arglines[i].substr(1, arglines[i].length()-2);
		else 
			arglinesProcessed[i] = arglines[i];

	}
#if defined(_SSA_DEBUG_TESTING_)
	MessageBox(0, L"Checking argument lines OK", LAPP_CAPTION, MB_OK);
#endif


	// Check for View is Running and wait for close
	bool bIsRunning = false;
	if ( arglinesProcessed[4] == L"1")
	{
#if defined(_SSA_DEBUG_TESTING_)
		MessageBox(0, L"Start Stage 1 (Checking Process)", LAPP_CAPTION, MB_OK);
#endif
		bool bContinueWait = false;
		do {
			bIsRunning = false;
			bContinueWait = false;
			int waitCountEnd = 30; // 30 sec
			int waitCount = 0;
			do {
//				if   ( arglinesProcessed[3] != L"0" )
//				{
					bIsRunning = CheckProcessID(arglinesProcessed[3], arglinesProcessed[0]);
//				}
//				else
//				{
//#ifndef _DEBUG
//					SingleInstance dcapp(_T("{VIEWDC-AEE8350A-B49A-4753-AB4B-E55479A48351}"));
//#else
//					SingleInstance dcapp(_T("{VIEWDC-AEE8350A-B49A-4753-AB4B-E55479A48350}"));
//#endif
//					bIsRunning  = ( dcapp.IsAnotherInstanceRunning() == TRUE);
//				}
				if ( bIsRunning  ){
					waitCount++;
					::Sleep(1000);
				}
				if (waitCount == waitCountEnd)
					break;

			}while ( bIsRunning );

			if (bIsRunning) 
			{
				// While waitCount == waitCountEnd
				// Show Warning Message
				bContinueWait =  ( MessageBox(0, LAPP_STILL_EXIST, LAPP_CAPTION, MB_YESNO | MB_ICONERROR) == IDYES);
			}

		}while (bContinueWait);
#if defined(_SSA_DEBUG_TESTING_)
		MessageBox(0, L"End Stage 1 (Process not loaded)", LAPP_CAPTION, MB_OK);
#endif
	}

	bool isFailedUpdate = false;
	bool isNeedRestart = false;
	wstring restartParameter = L"";
	if (!bIsRunning && (  arglinesProcessed[4] == L"1" || arglinesProcessed[4] == L"2" ) )
	{
#if defined(_SSA_DEBUG_TESTING_)
		MessageBox(0, L"Start Stage 2 - Try to Update", LAPP_CAPTION, MB_OK);
#endif
		// Checking Update folder and try to Update if Vista and Failed -> Restart with UAC
		size_t pos = arglinesProcessed[0].find_last_of(L"\\");
		wstring sourceFolder = arglinesProcessed[0];
		if (pos != wstring::npos)
			sourceFolder = arglinesProcessed[0].substr(0, pos);

#if defined(_SSA_DEBUG_TESTING_)
		MessageBox(0, L"Stage 2 - Load Update Dialog", LAPP_CAPTION, MB_OK);
#endif
		BOOL isFinalTry = (arglinesProcessed[4] == L"2" || isVista == FALSE  || ( isVista == TRUE && isElevated == TRUE));

		UpdateDlg* dlg = new UpdateDlg(sourceFolder, arglinesProcessed[1], arglinesProcessed[4], 0, isFinalTry);
		DWORD res = dlg->DoModal(0);

		if (res == IDABORT)
		{
			if (arglinesProcessed[4] == L"2" ||
				isVista == FALSE  || ( isVista == TRUE && isElevated == TRUE))
			{
				DWORD err = dlg->GetError();
				wchar_t buff[512];
				_snwprintf_s(buff, _countof(buff), _countof(buff)-1, LAPP_ERROR, err);
				// Failed to install
				MessageBox(0, buff, LAPP_CAPTION, MB_OK | MB_ICONERROR);
				isFailedUpdate = true;
			}else
			{
				// Restart Application as Elevated with 2
				isNeedRestart = true;
				restartParameter = L"2";
			}
		}else
			if (res == IDCANCEL)
			{
				// User Cancel
				MessageBox(0, LUSER_CANCELED, LAPP_CAPTION, MB_OK | MB_ICONERROR);
				isFailedUpdate = true;
			}else
				if (res == IDOK)
				{
					if (arglinesProcessed[4] == L"2")
					{
						isNeedRestart = true;
						restartParameter = L"3";
					}
				}
#if defined(_SSA_DEBUG_TESTING_)
		MessageBox(0, L"Finishing Stage 2 - Load Update Dialog", LAPP_CAPTION, MB_OK);
#endif
	}
	if (isNeedRestart)
	{
#if defined(_SSA_DEBUG_TESTING_)
		MessageBox(0, L"Restart ViewUpdate for UAC ", LAPP_CAPTION, MB_OK);
#endif
		TCHAR szPath[ MAX_PATH ];

		::GetModuleFileName( NULL, szPath, _tsizeof(szPath) );
		wstring params = arglines[0]+L" "+arglines[1]+L" "+arglines[2]+L" "+arglines[3]+L" "+restartParameter;

		if (restartParameter == L"2")
		{
			// Run evaluated
			RunElevated(NULL, szPath, params.c_str());
		}else
			if (restartParameter == L"3")
			{
				// Run not-evaluated
				RunAsStdUser(NULL, szPath, params.c_str());
			}
	}else
		if (!bIsRunning && !isFailedUpdate && !isNeedRestart)
		{
#if defined(_SSA_DEBUG_TESTING_)
		MessageBox(0, L"Stage 3 - Deleting update Folder ", LAPP_CAPTION, MB_OK);
#endif
			// Delete Update folder and
			size_t pos = arglinesProcessed[0].find_last_of(L"\\");
			wstring sourceFolder = arglinesProcessed[0];
			if (pos != wstring::npos)
				sourceFolder = arglinesProcessed[0].substr(0, pos);
#if defined(_SSA_DEBUG_TESTING_)
		MessageBox(0, L"Stage 3 - Start delete dialog", LAPP_CAPTION, MB_OK);
#endif
			UpdateDlg* dlg = new UpdateDlg(sourceFolder, arglinesProcessed[1], arglinesProcessed[4], TRUE, TRUE);
			DWORD res = dlg->DoModal(0);

			// Run ViewlinkDC++ again
			// MessageBox(0, CTSTRING(VIEWUPDATE_SUCCESS), CTSTRING(VIEWUPDATE), MB_OK | MB_ICONEXCLAMATION);
			if (arglinesProcessed[2] == L"1") // Start application
			{
				if (isVista)
				{
					RunAsStdUser(NULL, arglines[0].c_str());
				}else
				{
					MyShellExec(NULL, _T("open"), arglines[0].c_str());
				}
			}else
			{
				MessageBox(0, LUPDATA_SUCCESS, LAPP_CAPTION, MB_OK);
			}
#if defined(_SSA_DEBUG_TESTING_)
		MessageBox(0, L"End Stage 3 - Deleting update Folder ", LAPP_CAPTION, MB_OK);
#endif
		}

		::CoUninitialize();

		return 1;

}


